import logging
import uuid
from datetime import datetime
from typing import Dict, Any
import threading
import traceback

logger = logging.getLogger(__name__)

class TaskManager:
    def __init__(self):
        self.tasks = {}
        self.report_workflow = None
        logger.info("TaskManager 初始化完成")
        # 延迟初始化 workflow 避免循环导入
        self._init_workflow()

    def _init_workflow(self):
        """延迟初始化工作流"""
        try:
            from agent.reportAgent.workflow import ReportGenerationWorkflow
            self.report_workflow = ReportGenerationWorkflow()
            logger.info("ReportGenerationWorkflow 初始化完成")
        except Exception as e:
            logger.error(f"初始化工作流失败: {str(e)}")
            self.report_workflow = None

    def create_task(self, requirement: str, task_type: str = "report_generation", 
                   priority: str = "normal") -> Dict[str, Any]:
        """创建任务"""
        try:
            task_id = str(uuid.uuid4())
            task = {
                'id': task_id,
                'requirement': requirement,
                'task_type': task_type,
                'priority': priority,
                'status': 'pending',
                'created_at': datetime.now(),
                'agent_status': {},
                'result': None,
                'error': None,
                'logs': [],
                'progress': 0
            }
            
            self.tasks[task_id] = task
            
            # 使用线程处理任务
            thread = threading.Thread(target=self.process_task, args=(task_id,))
            thread.daemon = True
            thread.start()
            
            return task
            
        except Exception as e:
            logger.error(f"创建任务失败: {str(e)}")
            raise

    def process_task(self, task_id: str):
        """处理任务的方法（同步）"""
        try:
            task = self.tasks[task_id]
            task['status'] = 'processing'
            
            # 如果工作流未初始化，尝试重新初始化
            if not self.report_workflow:
                self._init_workflow()
            
            if not self.report_workflow:
                raise Exception("ReportGenerationWorkflow 未能正确初始化")
            
            # 设置当前任务ID到workflow，以便它能更新状态
            self.report_workflow.current_task_id = task_id
            self.report_workflow.task_manager = self
            
            # 调用 report_workflow 生成报告（同步调用）
            result = self.report_workflow.generate_report(
                requirement=task['requirement'],
                task_type=task['task_type'],
                priority=task['priority']
            )
            
            logger.info(f"工作流返回结果: {result}")
            
            if result['success']:
                task['status'] = 'completed'
                task['result'] = result
                
                # 更新任务的各种数据
                if 'agent_status' in result:
                    task['agent_status'] = result['agent_status']
                if 'optimized_prompt' in result:
                    task['optimized_prompt'] = result['optimized_prompt']
                if 'plan' in result:
                    task['plan'] = result['plan']
                if 'content' in result:
                    task['content'] = result['content']
                if 'quality_score' in result:
                    task['quality_score'] = result['quality_score']
                if 'logs' in result:
                    task['logs'] = result['logs']
                    
                task['progress'] = 100
            else:
                task['status'] = 'failed'
                task['error'] = result.get('error')
                if 'agent_status' in result:
                    task['agent_status'] = result['agent_status']
                if 'errors' in result:
                    task['errors'] = result['errors']
                
        except Exception as e:
            logger.error(f"处理任务失败: {str(e)}")
            logger.error(f"错误详情: {traceback.format_exc()}")
            self.tasks[task_id]['status'] = 'failed'
            self.tasks[task_id]['error'] = str(e)

    def get_task_status(self, task_id: str) -> Dict[str, Any]:
        """获取任务详细状态"""
        if task_id not in self.tasks:
            return {"error": "任务不存在"}
        
        task = self.tasks[task_id]
        
        # 构建详细状态信息
        detailed_status = {
            "agent_status": task.get("agent_status", {}),
            "logs": task.get("logs", []),
            "errors": task.get("errors", []),
            "progress": task.get("progress", 0)
        }
        
        # 添加规划代理的结果
        if "optimized_prompt" in task:
            detailed_status["optimized_prompt"] = task["optimized_prompt"]
            
        if "plan" in task:
            detailed_status["plan"] = task["plan"]
            
        # 添加其他代理的结果
        if "content" in task:
            detailed_status["content"] = task["content"]
            
        if "quality_score" in task:
            detailed_status["quality_score"] = task["quality_score"]
        
        return detailed_status

    def update_task_partial_status(self, task_id: str, partial_data: Dict[str, Any]):
        """部分更新任务状态（用于实时更新）"""
        if task_id not in self.tasks:
            return False
            
        task = self.tasks[task_id]
        
        # 更新指定的字段
        for key, value in partial_data.items():
            if value is not None:
                task[key] = value
        
        logger.info(f"任务 {task_id} 部分状态已更新: {partial_data}")
        return True

# 创建全局实例
task_manager = TaskManager()

def get_task_manager():
    """获取任务管理器实例"""
    return task_manager 